package org.geomesa.example.hbase;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.hbase.*;
import org.apache.hadoop.hbase.client.*;
import org.apache.hadoop.hbase.util.Bytes;
import org.geotools.data.DataStore;
import org.geotools.data.DataStoreFinder;
import org.geotools.data.Query;
import org.geotools.data.simple.SimpleFeatureIterator;
import org.geotools.filter.text.ecql.ECQL;

import java.io.IOException;
import java.util.*;

public class GridImageQueryService {

    private final String zkQuorum;
    private final String geomesaCatalog;
    private final String invertIndexTableName;
    private final String forwardIndexTableName;

    private static final String INVERT_CF = "img";
    private static final String FORWARD_CF = "grid";

    private Connection connection;

    public GridImageQueryService(String zkQuorum, String geomesaCatalog, String invertIndexTableName, String forwardIndexTableName) {
        this.zkQuorum = zkQuorum;
        this.geomesaCatalog = geomesaCatalog;
        this.invertIndexTableName = invertIndexTableName;
        this.forwardIndexTableName = forwardIndexTableName;
    }

    public void init() throws IOException {
        Configuration config = HBaseConfiguration.create();
        config.set("hbase.zookeeper.quorum", zkQuorum);
        connection = ConnectionFactory.createConnection(config);
    }

    public void close() throws IOException {
        if (connection != null) {
            connection.close();
        }
    }

    /**
     * 主查询方法：输入空间范围，输出每张影像中覆盖的网格及其window位置
     */
    public Map<String, Map<String, String>> query(double minLon, double minLat, double maxLon, double maxLat) throws IOException {
        Map<String, Map<String, String>> result = new HashMap<>();

        // 1. 查空间范围包含哪些网格（使用GeoMesa）
        List<String> gridKeys = getCoveredGridKeys(minLon, minLat, maxLon, maxLat);
        System.out.println("空间范围包含 " + gridKeys.size() + " 个网格");

        // 2. 倒排索引查影像
        Map<String, List<String>> imageToGrids = new HashMap<>();
        Table invertTable = connection.getTable(TableName.valueOf(invertIndexTableName));
        for (String gridKey : gridKeys) {
            Get get = new Get(Bytes.toBytes(gridKey));
            Result res = invertTable.get(get);
            for (Cell cell : res.rawCells()) {
                String column = Bytes.toString(CellUtil.cloneQualifier(cell)); // img:img_1
                String imgId = column.replace("img_", "");
                imageToGrids.computeIfAbsent(imgId, k -> new ArrayList<>()).add(gridKey);
            }
        }
        invertTable.close();

        // 3. 正向索引查每个网格的window
        Table forwardTable = connection.getTable(TableName.valueOf(forwardIndexTableName));
        for (Map.Entry<String, List<String>> entry : imageToGrids.entrySet()) {
            String imgId = entry.getKey();
            List<String> grids = entry.getValue();

            Get get = new Get(Bytes.toBytes("img_" + imgId));
            Result res = forwardTable.get(get);
            Map<String, String> gridToWindow = new HashMap<>();
            for (String gridKey : grids) {
                String column = "grid:" + gridKey;
                byte[] value = res.getValue(Bytes.toBytes(FORWARD_CF), Bytes.toBytes(column));
                gridToWindow.put(gridKey, value != null ? Bytes.toString(value) : "[unknown]");
            }
            result.put(imgId, gridToWindow);
        }
        forwardTable.close();

        return result;
    }

    /**
     * 查询空间范围内包含的 gridKey（调用GeoMesa 查询空间索引）
     */
    private List<String> getCoveredGridKeys(double minLon, double minLat, double maxLon, double maxLat) {
        List<String> result = new ArrayList<>();
        DataStore dataStore = null;

        try {
            Map<String, String> params = new HashMap<>();
            params.put("hbase.zookeepers", zkQuorum);
            params.put("hbase.catalog", geomesaCatalog);

            dataStore = DataStoreFinder.getDataStore(params);
            if (dataStore == null) throw new RuntimeException("无法连接GeoMesa DataStore");

            String cql = String.format("BBOX(geom, %f, %f, %f, %f)", minLon, minLat, maxLon, maxLat);
            Query query = new Query("custom_polygon", ECQL.toFilter(cql), new String[]{"name"});

            try (SimpleFeatureIterator it = dataStore.getFeatureSource("custom_polygon").getFeatures(query).features()) {
                while (it.hasNext()) {
                    result.add((String) it.next().getAttribute("name"));
                }
            }
        } catch (Exception e) {
            throw new RuntimeException("查询GeoMesa网格失败", e);
        } finally {
            if (dataStore != null) {
                dataStore.dispose();
            }
        }

        return result;
    }

    public static void main(String[] args) throws IOException {
        GridImageQueryService queryService = new GridImageQueryService(
                "localhost:2181",
                "globalgrid2_8",
                "globalgrid2_8_invert_index",
                "globalgrid2_8_forward_index"
        );

        queryService.init();

        // 输入一个查询区域
        double minLon = 138.0, minLat = 25.0, maxLon = 140.0, maxLat = 30.0;

        Map<String, Map<String, String>> result = queryService.query(minLon, minLat, maxLon, maxLat);

        for (String imgId : result.keySet()) {
            System.out.println("影像：" + imgId);
            Map<String, String> gridToWindow = result.get(imgId);
            for (Map.Entry<String, String> e : gridToWindow.entrySet()) {
                System.out.println("  网格: " + e.getKey() + " -> 像素窗口: " + e.getValue());
            }
        }

        queryService.close();
    }
}
